 # -*- coding: utf-8 -*-

import math as m
import random as r
import numpy as np
#import argparse 

from abaqus import *
from abaqusConstants import *
from caeModules import *
from driverUtils import executeOnCaeStartup

#parser = argparse.ArgumentParser()
#parser.add_argument("--output_path", default='/workspace/outputs/possion/ce/', type=str, help='output path')

#args = parser.parse_args()
#print(args.output_path)

output_path='/tmp/'

def poisson(code):
    Mdb()
    num=0.5*len(code)+1.0
    all_num=int(2.0*num)
    size=0.5
    load_magnitude=0.2
    mesh_size=0.25

    #produce all cube
    for i in range(1,int(all_num)):
        s = mdb.models['Model-1'].ConstrainedSketch(name='__profile__', 
                                                    sheetSize=200.0)
        s.setPrimaryObject(option=STANDALONE)
        s.rectangle(point1=(size, size), 
                    point2=(0.0, 0.0))
        p = mdb.models['Model-1'].Part(name='Part-%s'%i, 
                                       dimensionality=THREE_D, 
                                       type=DEFORMABLE_BODY)
        p = mdb.models['Model-1'].parts['Part-%s'%i]
        p.BaseSolidExtrude(sketch=s, 
                           depth=size)
        s.unsetPrimaryObject()
        del mdb.models['Model-1'].sketches['__profile__']
    
    #all cube enter assemble    
    a = mdb.models['Model-1'].rootAssembly
    a.DatumCsysByDefault(CARTESIAN)
    for i in range(1,int(all_num)):
        a.Instance(name='Part-%s-1'%i, 
                   part=mdb.models['Model-1'].parts['Part-%s'%i], 
                   dependent=ON)


    direction=[]
    for i in range(0,len(code)):    
        if code[i]==0:
            direction.append([size,0,0])
        if code[i]==1:
            direction.append([0,size,0])
    #print(direction)
    
    for j in range(2,int(all_num)):
        for i in range(j,int(all_num)):
            mdb.models['Model-1'].rootAssembly.translate\
            (instanceList=('Part-%s-1'%i, ), 
             vector=direction[j-2])
    
    assembly=''
    a1 = mdb.models['Model-1'].rootAssembly
    for i in range(1,all_num):
        assemble_parts="a1.instances['"+'Part-%s-1'%i+"'],"
        assembly=assembly+assemble_parts
    
    a1.InstanceFromBooleanMerge(name='Part-%s'%all_num, 
                                instances=(tuple(eval(assembly))), 
                                originalInstances=SUPPRESS,
                                domain=GEOMETRY)
    
    s = mdb.models['Model-1'].ConstrainedSketch(name='__profile__', 
                                                sheetSize=200.0)
    s.setPrimaryObject(option=STANDALONE)
    s.rectangle(point1=(num*size, 
                        num*size), 
                point2=(-(num-4.5)*size, 
                        (num-1)*size))
    p = mdb.models['Model-1'].Part(name='Part-%s'%(all_num+1), 
                                   dimensionality=THREE_D, 
                                   type=DEFORMABLE_BODY)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+1)]
    p.BaseSolidExtrude(sketch=s, 
                       depth=size)
    s.unsetPrimaryObject()
    del mdb.models['Model-1'].sketches['__profile__']
    
    s = mdb.models['Model-1'].ConstrainedSketch(name='__profile__', 
                                                sheetSize=200.0)
    s.setPrimaryObject(option=STANDALONE)
    s.rectangle(point1=(num*size, 
                        num*size), 
                point2=((num-1)*size, 
                        (num+1)*size))
    p = mdb.models['Model-1'].Part(name='Part-%s'%(all_num+2), 
                                   dimensionality=THREE_D, 
                                   type=DEFORMABLE_BODY)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+2)]
    p.BaseSolidExtrude(sketch=s, 
                       depth=size)
    s.unsetPrimaryObject()
    del mdb.models['Model-1'].sketches['__profile__']
    
    a = mdb.models['Model-1'].rootAssembly
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+1)]
    a.Instance(name='Part-%s-1'%(all_num+1), 
               part=p, 
               dependent=ON)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+2)]
    a.Instance(name='Part-%s-1'%(all_num+2), 
               part=p, 
               dependent=ON)
    
    mdb.models['Model-1'].rootAssembly.RadialInstancePattern\
    (instanceList=('Part-%s-1'%all_num, 
                   'Part-%s-1'%(all_num+1),
                   'Part-%s-1'%(all_num+2)), 
    point=(0.0, 0.5*size, 0.5*size), 
    axis=(1.0, 0.0, 0.0), 
    number=4, 
    totalAngle=360.0)
    
    a1 = mdb.models['Model-1'].rootAssembly
    a1.InstanceFromBooleanMerge(name='Part-%s'%(all_num+3), instances=(
        a1.instances['Part-%s-1'%all_num], a1.instances['Part-%s-1'%(all_num+1)], 
        a1.instances['Part-%s-1'%(all_num+2)], a1.instances['Part-%s-1-rad-2'%all_num], 
        a1.instances['Part-%s-1-rad-3'%all_num], a1.instances['Part-%s-1-rad-4'%all_num], 
        a1.instances['Part-%s-1-rad-2'%(all_num+1)], a1.instances['Part-%s-1-rad-3'%(all_num+1)], 
        a1.instances['Part-%s-1-rad-4'%(all_num+1)], a1.instances['Part-%s-1-rad-2'%(all_num+2)], 
        a1.instances['Part-%s-1-rad-3'%(all_num+2)], a1.instances['Part-%s-1-rad-4'%(all_num+2)], ), 
        originalInstances=SUPPRESS, domain=GEOMETRY)
    
    a1 = mdb.models['Model-1'].rootAssembly
    a1.RadialInstancePattern(instanceList=('Part-%s-1'%(all_num+3), ), 
                             point=(-0.5*size, 0.5*size, 0.0), 
                             axis=(0.0, 0.0, 1.0), 
                             number=2, 
                             totalAngle=360.0)
    
    a1 = mdb.models['Model-1'].rootAssembly
    a1.InstanceFromBooleanMerge(name='Part-%s'%(all_num+4), instances=(
        a1.instances['Part-%s-1'%(all_num+3)], 
        a1.instances['Part-%s-1-rad-2'%(all_num+3)], ), 
        originalInstances=SUPPRESS, 
        domain=GEOMETRY)
    
    s1 = mdb.models['Model-1'].ConstrainedSketch(name='__profile__', 
                                                 sheetSize=200.0)
    s1.setPrimaryObject(option=STANDALONE)
    s1.rectangle(point1=(size, 0.0),
                 point2=((num+1)*size, size))
    p = mdb.models['Model-1'].Part(name='Part-%s'%(all_num+5), 
                  dimensionality=THREE_D, 
                  type=DEFORMABLE_BODY)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+5)]
    p.BaseSolidExtrude(sketch=s1,
                       depth=size)
    s1.unsetPrimaryObject()
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+5)]
    del mdb.models['Model-1'].sketches['__profile__']
    
    s = mdb.models['Model-1'].ConstrainedSketch(name='__profile__', 
                                                sheetSize=200.0)
    s.setPrimaryObject(option=STANDALONE)
    s.rectangle(point1=(-2*size, size), 
                point2=(-(num+2)*size, 0.0))
    p = mdb.models['Model-1'].Part(name='Part-%s'%(all_num+6), 
                  dimensionality=THREE_D, 
                  type=DEFORMABLE_BODY)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+6)]
    p.BaseSolidExtrude(sketch=s, 
                       depth=size)
    s.unsetPrimaryObject()
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+6)]
    
    a1 = mdb.models['Model-1'].rootAssembly
    del mdb.models['Model-1'].sketches['__profile__']
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+5)]
    a1.Instance(name='Part-%s-1'%(all_num+5), 
                part=p, 
                dependent=ON)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+6)]
    a1.Instance(name='Part-%s-1'%(all_num+6), 
                part=p, 
                dependent=ON)
    
    a1 = mdb.models['Model-1'].rootAssembly
    a1.InstanceFromBooleanMerge(name='Part-%s'%(all_num+7), 
                                instances=(a1.instances['Part-%s-1'%(all_num+4)], 
                                           a1.instances['Part-%s-1'%(all_num+5)], 
                                           a1.instances['Part-%s-1'%(all_num+6)], ), 
                                                        originalInstances=SUPPRESS, 
                                                        domain=GEOMETRY)
    
    #material
    mdb.models['Model-1'].Material(name='Material-1')
    mdb.models['Model-1'].materials['Material-1'].Elastic(table=((2200.0,
                                                                  0.394), ))
    mdb.models['Model-1'].HomogeneousSolidSection(name='Section-1', 
                                                  material='Material-1', 
                                                  thickness=None)
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+7)]
    region = p.Set(cells=p.cells[0:1], 
                   name='Set-1')
    
    p.SectionAssignment(region=region, 
                        sectionName='Section-1', 
                        offset=0.0, 
                        offsetType=MIDDLE_SURFACE, 
                        offsetField='', 
                        thicknessAssignment=FROM_SECTION)
    
    #step
    mdb.models['Model-1'].StaticStep(name='Step-1', 
                                     previous='Initial', 
                                     maxNumInc=10000, 
                                     stabilizationMagnitude=0.0002, 
                                     stabilizationMethod=DAMPING_FACTOR, 
                                     continueDampingFactors=False, 
                                     adaptiveDampingRatio=None, 
                                     initialInc=0.01, 
                                     nlgeom=ON)
    mdb.models['Model-1'].fieldOutputRequests['F-Output-1'].\
    setValues(variables=('S', 'U', 'COORD'))
    
    #load
    a = mdb.models['Model-1'].rootAssembly
    s1 = a.instances['Part-%s-1'%(all_num+7)].faces
    side1Faces1 = s1.findAt(((-(num+2)*size, 0.5*size, 0.5*size), ), 
                            (((num+1)*size, 0.5*size, 0.5*size), ))
    region = a.Surface(side1Faces=side1Faces1, 
                       name='Surf-1')
    mdb.models['Model-1'].Pressure(name='Load-1', 
                                   createStepName='Step-1', 
                                   region=region, 
                                   distributionType=UNIFORM, 
                                   field='', 
                                   magnitude=-load_magnitude, 
                                   amplitude=UNSET)
    
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+7)]
    p.DatumPlaneByPrincipalPlane(principalPlane=YZPLANE, 
                                 offset=(num-0.5)*size)
    c = p.cells
    pickedCells = c[0:1]
    p.PartitionCellByDatumPlane(datumPlane=p.datums[3], 
                                cells=pickedCells)
    
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+7)]
    p.DatumPlaneByPrincipalPlane(principalPlane=XYPLANE, 
                                 offset=0.5*size)
    c = p.cells
    pickedCells = c[0:1]+c[2:3]+c[4:6]
    p.PartitionCellByDatumPlane(datumPlane=p.datums[5], 
                                cells=pickedCells)
    
    #mesh
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+7)]
    p.seedPart(size=mesh_size, 
               deviationFactor=0.1, 
               minSizeFactor=0.1)
    c = p.cells
    pickedRegions = c.getSequenceFromMask(mask=('[#ff ]', ), )
    p.setMeshControls(regions=pickedRegions, 
                      elemShape=TET, 
                      technique=FREE)
    elemType1 = mesh.ElemType(elemCode=C3D20R)
    elemType2 = mesh.ElemType(elemCode=C3D15)
    elemType3 = mesh.ElemType(elemCode=C3D10)
    c = p.cells
    cells = c.getSequenceFromMask(mask=('[#ff ]', ), )
    pickedRegions =(cells, )
    p.setElementType(regions=pickedRegions, 
                     elemTypes=(elemType1, elemType2, elemType3))
    p.generateMesh()
    
    #set
    p.Set(vertices=p.vertices.findAt((((num-0.5)*size, -num*size, 0.5*size), ), 
                                     (((num-0.5)*size, (num+1)*size, 0.5*size), )), \
          name='Set-2')   
    p.Set(vertices=p.vertices.findAt(((-(num+2)*size, 0, size), ),
                                     (((num+1)*size, 0, size), )), 
          name='Set-3')   
    
    #job
    mdb.Job(name='Job-1', model='Model-1', description='', type=ANALYSIS, 
        atTime=None, waitMinutes=0, waitHours=0, queue=None, memory=90, 
        memoryUnits=PERCENTAGE, getMemoryFromAnalysis=True, 
        explicitPrecision=SINGLE, nodalOutputPrecision=SINGLE, echoPrint=OFF, 
        modelPrint=OFF, contactPrint=OFF, historyPrint=OFF, userSubroutine='', 
        scratch='', resultsFormat=ODB, multiprocessingMode=DEFAULT, numCpus=1, 
        numGPUs=0)
    mdb.jobs['Job-1'].submit(consistencyChecking=OFF)
    mdb.jobs['Job-1'].waitForCompletion()
    
    #post processing
    #y_direction
    session.viewports['Viewport: 1'].setValues(displayedObject=session.openOdb\
                     (name='/workspace/Job-1.odb'))
    vp = session.viewports[session.currentViewportName]
    odb = vp.displayedObject
    lastFrame=odb.steps['Step-1'].frames[-1]
    displacement=lastFrame.fieldOutputs['COORD']
    center=odb.rootAssembly.instances['PART-%s-1'%(all_num+7)].nodeSets['SET-2']
    centerDisplacement=displacement.getSubset(region=center)
    centerValues=centerDisplacement.values
    y_deformation=m.sqrt(m.pow(centerValues[1].data[0]-centerValues[0].data[0],2)+
                      m.pow(centerValues[1].data[1]-centerValues[0].data[1],2)+
                      m.pow(centerValues[1].data[2]-centerValues[0].data[2],2))
    y_strain=(y_deformation-(2*num+1)*size)/((2*num+1)*size)
#    print('y_strain')    print(y_strain)
    
    #x_direction
    session.viewports['Viewport: 1'].setValues(displayedObject=session.openOdb\
                     (name='/workspace/Job-1.odb'))
    vp = session.viewports[session.currentViewportName]
    odb = vp.displayedObject
    lastFrame=odb.steps['Step-1'].frames[-1]
    displacement=lastFrame.fieldOutputs['COORD']
    center=odb.rootAssembly.instances['PART-%s-1'%(all_num+7)].nodeSets['SET-3']
    centerDisplacement=displacement.getSubset(region=center)
    centerValues=centerDisplacement.values 
    x_deformation=m.sqrt(m.pow(centerValues[1].data[0]-centerValues[0].data[0],2)+
                      m.pow(centerValues[1].data[1]-centerValues[0].data[1],2)+
                      m.pow(centerValues[1].data[2]-centerValues[0].data[2],2))
    x_strain=(x_deformation-(2*num+3)*size)/((2*num+3)*size)
#    print('x_strain')    print(x_strain)
    poisson_ratio=(y_strain/x_strain)
    f.write(str(code)+'\t')
    f.write(str(poisson_ratio)+'\t')
    
    
    p = mdb.models['Model-1'].parts['Part-%s'%(all_num+7)]
    volume=p.getVolume()
    surface=p.getArea(p.faces)
    
    dic[str(code)]=poisson_ratio,volume,surface
    f.write(str(volume)+'\t')
    f.write(str(surface)+'\n')
    
    return poisson_ratio

def get_fitness(pop):
    pred=pop*np.mat(np.ones([DNA_SIZE,1])) 
    pred=np.copy(pred.T)
    for i in range(0,POP_SIZE): 
        if pred[:,i]==0.5*DNA_SIZE:
            if str(pop[i,:]) in dic.keys():
                pred[:,i]=dic[str(pop[i,:])][0]
                f.write(str(pop[i,:])+'\t')
                f.write(str(dic[str(pop[i,:])][0])+'\t')
                f.write(str(dic[str(pop[i,:])][1])+'\t')
                f.write(str(dic[str(pop[i,:])][2])+'\n')
            else:
                pred[:,i]=poisson(pop[i,:])
        else:
            pred[:,i]=0
            f.write(str(pred[:,i])+'\n')
            
    fitness=np.array((pred-np.min(pred))+1e-4)
    return fitness[0]

def crossover_and_mutation(pop):
    new_pop = []  
    for j in range(0,POP_SIZE):
        child = np.copy(pop[j])
        if np.random.rand() < 0.7:
            k=np.copy(j)           
            while k==j:
                k=np.random.randint(POP_SIZE)
                if k!=j:
                    break
            mother=np.copy(pop[k])  
            cross_points = np.random.randint(low=0, 
                                             high=DNA_SIZE)
            child[cross_points:] = np.copy(mother[cross_points:])
        mutation(child)
        new_pop.append(child)
    return new_pop

def mutation(child):
    if np.random.rand()<0.2: 
        mutate_point = np.random.randint(0,DNA_SIZE)
        child[mutate_point] = child[mutate_point]^1
        mutate_point = np.random.randint(0,DNA_SIZE)
        child[mutate_point] = child[mutate_point]^1
    else:
        child=child
    return child

def select(pop, fitness):
    idx=np.random.choice(np.arange(POP_SIZE), 
                         size=POP_SIZE, 
                         replace=True,
                         p=(fitness)/(fitness.sum()))
    return pop[idx]


if __name__ == "__main__":
    DNA_SIZE=16
    POP_SIZE=60
    N_GENERATIONS=300
    pop=[]
    dic={}
    
    for _ in range(POP_SIZE):
        one_code=[]
        for _ in range(int(0.5*DNA_SIZE)):
            one_code.append(1)
            one_code.append(0)
        r.shuffle(one_code) 
        pop.append(one_code)  
    pop=np.array(pop) 
    
    
    
    for i in range(0,N_GENERATIONS):
        #print(pop)
        print('*'*15+"%s"%(i+1)+'*'*15)
        f=open(output_path+'record.txt','a')
        f.write(str("*"*15+"%s"%str(i+1.0)+"*"*15+'\n'))       
#        f.write(str(pop)+'\n')
        fitness = get_fitness(pop)        
        pop = select(pop, fitness)
        pop = np.array(crossover_and_mutation(pop))
        f.close() 
    
c=list(dic.keys())
g=list(dic.values())
for i in range(0,len(c)):
    f=open(output_path+'code_list.txt','a')
    f.write(str(c[i])+'\n')
    f.close()
    p=open(output_path+'possion_list.txt','a')
    p.write(str(g[i])+'\n')
    p.close()
